﻿using System.Collections;
using System.Collections.Generic;
using System.Windows.Controls;

namespace TfsBranchLib2.Model
{
    public class TreeViewItemsEnumerable : IEnumerable<TreeViewItem>
    {
        private readonly TreeViewItem item;

        public TreeViewItemsEnumerable(TreeViewItem item)
        {
            this.item = item;
        }

        public IEnumerator<TreeViewItem> GetEnumerator()
        {
            return new TreeViewEnumerator(item);
        }

        IEnumerator IEnumerable.GetEnumerator()
        {
            return new TreeViewEnumerator(item);
        }
    }

    public sealed class TreeViewEnumerator : IEnumerator<TreeViewItem>, IEnumerator
    {
        private readonly TreeViewItem baseElement;

        private TreeViewItem currentElement;

        private IEnumerator currentEnumerator;

        private readonly Stack<IEnumerator> stack = new Stack<IEnumerator>();

        public TreeViewEnumerator(TreeViewItem baseItem)
        {
            baseElement = baseItem;
        }

        public object Current
        {
            get { return currentElement; }
        }

        public bool MoveNext()
        {
            if (currentElement == null)
            {
                if (baseElement == null)
                    return false;
                currentElement = baseElement;
                return true;
            }
            if (currentEnumerator == null)
            {
                currentEnumerator = ((IEnumerable) (currentElement.Items)).GetEnumerator();
            }
            if (currentEnumerator.MoveNext())
            {
                SetCurrentElement();
                return true;
            }
            while (stack.Count > 0)
            {
                currentEnumerator = stack.Pop();
                if (currentEnumerator.MoveNext())
                {
                    SetCurrentElement();
                    return true;
                }
            }
            return false;
        }

        public void Reset()
        {
            stack.Clear();
            currentElement = null;
            currentEnumerator = null;
        }

        public void Dispose()
        {
            Reset();
        }

        TreeViewItem IEnumerator<TreeViewItem>.Current
        {
            get { return currentElement; }
        }

        private void SetCurrentElement()
        {
            currentElement = (TreeViewItem) currentEnumerator.Current;
            if (currentElement.HasItems)
            {
                stack.Push(currentEnumerator);
                currentEnumerator = null;
            }
        }
    }
}